import argparse  # 导入argparse库，用于解析命令行参数
import datetime  # 导入datetime库，用于处理日期和时间
import json  # 导入json库，用于处理JSON数据
import numpy as np  # 导入numpy库，用于科学计算
import os  # 导入os库，用于与操作系统交互
import time  # 导入time库，用于处理时间相关功能
from pathlib import Path  # 从pathlib模块导入Path类，用于处理文件系统路径
import torch.nn as nn  # 导入torch.nn模块，用于构建神经网络

import torch  # 导入torch库，PyTorch的核心库
import torch.backends.cudnn as cudnn  # 导入cudnn模块，用于配置CUDA深度神经网络库
from torch.utils.tensorboard import SummaryWriter  # 从torch.utils.tensorboard导入SummaryWriter类，用于记录TensorBoard日志
import torchvision.transforms as transforms  # 导入torchvision.transforms模块，用于图像预处理
import torchvision.datasets as datasets  # 导入torchvision.datasets模块，用于加载常用数据集

import timm  # 导入timm库，用于加载预训练模型和优化器

# assert timm.__version__ == "0.3.2"  # 检查timm库的版本，确保为0.3.2（此行代码被注释掉了）
import timm.optim.optim_factory as optim_factory  # 导入timm的优化器工厂模块

import util.misc as misc  # 导入自定义的util.misc模块，包含一些实用函数
from util.misc import NativeScalerWithGradNormCount as NativeScaler  # 从util.misc导入NativeScalerWithGradNormCount类，并重命名为NativeScaler

import models_Mambamae  # 导入自定义的models_Mambamae模块，包含MedMamba模型的定义
from VMamba.classification.medical.medmamba import VSSM  # 从VMamba.classification.medical.medmamba导入VSSM类，用于加载Swin-UMamba模型
#from VMamba.classification.models.vmamba import VSSM

from engine_pretrain_vmamba import train_one_epoch_MedMamba  # 从engine_pretrain模块导入train_one_epoch_MedMamba函数，用于MedMamba模型的预训练

def load_medmamba_model(device='cuda'):
    """加载Swin-UMamba模型"""
    # 创建VSSM类的实例，初始化模型结构
    model = VSSM(depths=[2, 2, 4, 2], dims=[96, 192, 384, 768], num_classes=4).to("cuda")
    model = model.cuda()
    # 指定预训练权重文件的路径
    checkpoint_path = "/home/ma-user/work/mae/check_point/MedMamba.pth"
    # 加载预训练权重
    checkpoint = torch.load(checkpoint_path, map_location=device)
    # 将预训练权重加载到模型中
    model.load_state_dict(checkpoint, strict=False)  # 根据需要选择是否要 strict

    return model


# 加载Swin-UMamba模型
medmamba = load_medmamba_model()
medmamba.to(device)
medmamba.eval()  # 将模型设置为评估模式

# 3. 准备验证集
val_dataset = YourDatasetClass(root='path_to_val_data', train=False)  # 根据你的数据集调整
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)

# 4. 运行验证
correct = 0
total = 0
with torch.no_grad():  # 禁用梯度计算
    for data in val_loader:
        inputs, labels = data
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy on validation set: {100 * correct / total:.2f}%')